import torch
from torch import nn

torch.device('cpu')
torch.cuda.device('cuda')
torch.cuda.device('cuda:1')
print(torch.cuda.device_count())
print(torch.__version__)

def tryGPU(i=0):
    if torch.cuda.device_count() >= i+1:
        return torch.device(f'cuda:{i}')
    return torch.device('cpu')

def tryAllGPUs():
    devices = [torch.device(f'cuda:{i}') for i in range(torch.cuda.device_count())]
    return devices if devices else [torch.device('cpu')]

device = tryGPU()
devices = tryAllGPUs()
print(device)
print(devices)

x = torch.tensor([0,1,2],device=tryGPU())
y = torch.tensor([1,2,3])
y = y.cuda(0)
print(x+y)